import numpy as np
import cv2
import multiprocessing
cimport numpy as np
cimport cython
import math

from cython.parallel import prange

@cython.boundscheck(False)
@cython.wraparound(False)

# instruction:
# python3 setup.py build_ext --inplace

def propagate2(set rules_h, set rules_v,list wave_oriented, int Width, int Height):

    cdef int p1 = 0
    cdef int p2 = 0
    cdef int p3 = 0
    cdef int p4 = 0
    cdef int i,j,k
    cdef list w1
    cdef list w2
    for k in range(1):
        for j in range(Height):
            for i in range(Width):

                #if len(wave_oriented[j*Width+i]) != 1:
                #    continue

                i_0 = i-1
                j_0 = j-1
                i_1 = i+1
                j_1 = j+1
                wavelet_temp1 = set()
                wavelet_temp2 = set()
                wavelet_temp3 = set()
                wavelet_temp4 = set()

                # construct local context and look it up in rules
                w1 = wave_oriented[j*Width+i]
                # (i-1,j)
                if i_0 >= 0 and len(wave_oriented[j*Width+i_0]) != 1:
                    w2 = wave_oriented[j*Width+i_0]
                    for element_i in w1:
                        for element_j in w2:
                            if (element_j[0],element_j[1],element_i[0],element_i[1]) in rules_h:
                                wavelet_temp1.add(element_j)
                # (i+1,j)
                if i_1 < Width and len(wave_oriented[j*Width+i_1]) != 1:
                    w2 = wave_oriented[j*Width+i_1]
                    for element_i in w1:
                        for element_j in w2:
                            if (element_i[0],element_i[1],element_j[0],element_j[1]) in rules_h:
                                wavelet_temp2.add(element_j)
                # (i,j-1)
                if j_0 >= 0 and len(wave_oriented[j_0*Width+i]) != 1:
                    w2 = wave_oriented[j_0*Width+i]
                    for element_i in w1:
                        for element_j in w2:
                            if (element_j[0],element_j[1],element_i[0],element_i[1]) in rules_v:
                                wavelet_temp3.add(element_j)
                # (i,j+1)
                if j_1 < Height and len(wave_oriented[j_1*Width+i]) != 1:
                    w2 = wave_oriented[j_1*Width+i]
                    for element_i in w1:
                        for element_j in w2:
                            if (element_i[0],element_i[1],element_j[0],element_j[1]) in rules_v:
                                wavelet_temp4.add(element_j)

                # update neighbors
                if len(wavelet_temp1) != 0:
                    wave_oriented[j*Width+i_0].clear()
                    for element in wavelet_temp1:
                        wave_oriented[j*Width+i_0].append(element)
                if len(wavelet_temp2) != 0:
                    wave_oriented[j*Width+i_1].clear()
                    for element in wavelet_temp2:
                        wave_oriented[j*Width+i_1].append(element)
                if len(wavelet_temp3) != 0:
                    wave_oriented[j_0*Width+i].clear()
                    for element in wavelet_temp3:
                        wave_oriented[j_0*Width+i].append(element)
                if len(wavelet_temp4) != 0:
                    wave_oriented[j_1*Width+i].clear()
                    for element in wavelet_temp4:
                        wave_oriented[j_1*Width+i].append(element)


def get_weight_of_candidate_c(index, candidate, oriented_weight_dict):
    return oriented_weight_dict[index][candidate]

def shannon_entropy_of_wavelet_c(index,wavelet,oriented_weight_dict):

    entropy = 0
    sum_weight = 0
    probabilities = []
    for i in wavelet:
        probabilities.append(get_weight_of_candidate_c(index,i,oriented_weight_dict))

    sum_weight = sum(probabilities)
    probabilities = [x/sum_weight for x in probabilities]
    probabilities_arr = np.array(probabilities)

    entropy = (-probabilities_arr.dot(np.log(probabilities)))

    # add a small purturbation
    purturbation = np.random.rand() * 0.001
    # entropy += purturbation

    #print(purturbation)
    #input()

    return entropy

def update_wave_entropy_c(entropy,wave_oriented_old, wave_oriented, oriented_weight_dict):

    #if len(wave_oriented) == len(wave_oriented_old):
    #    for i in range(len(wave_oriented)):
    #        if wave_oriented_old[i] != wave_oriented[i]:
    #            entropy[i] = shannon_entropy_of_wavelet_c(i,wave_oriented[i],oriented_weight_dict)
    #else:

    for i in range(len(wave_oriented)):
        entropy[i] = shannon_entropy_of_wavelet_c(i,wave_oriented[i],oriented_weight_dict)


def spawn1(num,q):
    cdef int p = 0
    for i in range(num):
        p += 1
    q.put(p)

def spawn2(num):

    q = multiprocessing.Queue()

    cdef int total = 0
    cdef list p_list = []
    for i in range(4):
        p = multiprocessing.Process(target=spawn1, args=(num,q))
        p.start()
        p_list.append(p)

    for p in p_list:
        p.join()

    result = []
    for i in range(len(p_list)):
        result.append(q.get())

    return result
